#include <ctrl_alg/pid/pid_pos.hpp>
// #include <general_interface/msg/control_signal.hpp>
#include <general_interface/msg/chassis_servo.hpp>
#include <rclcpp/rclcpp.hpp>
#include <rclcpp_lifecycle/rclcpp_lifecycle/lifecycle_node.hpp>
#include <sensor_msgs/msg/imu.hpp>
#include <std_msgs/msg/float64.hpp>
#include <tf2/LinearMath/Matrix3x3.h>
#include <tf2/LinearMath/Quaternion.h>

static const int gyro_rate = 100;
static const int angle_rate = 50;
static const int speed_rate = 10;

class DataFilter
{
  private:
    double ewa;
    double alaph;
    double threshold;
    int count = 0;
    static const int Boost = 10;

  public:
    DataFilter(double alaph, double threshold)
        : ewa(0), alaph(alaph), threshold(threshold)
    {
    }
    double filter(double data)
    {
        count += 1;
        if (count > Boost)
        {
            double nv = ewa * alaph + data * (1 - alaph);
            if (std::abs(nv - ewa) < threshold)
                ewa = nv;
            return ewa;
        }
        return data;
    }
};

class ControlNode : public rclcpp_lifecycle::LifecycleNode
{
  private:
    using WheelSpeed = std_msgs::msg::Float64;
    using WheelTorque = std_msgs::msg::Float64;
    rclcpp_lifecycle::LifecyclePublisher<WheelTorque>::SharedPtr leftTorquePub,
        rightTorquePub;
    rclcpp::Subscription<WheelSpeed>::SharedPtr leftSpeedSub, rightSpeedSub;
    rclcpp::Subscription<sensor_msgs::msg::Imu>::SharedPtr imuSub;
    rclcpp::Subscription<general_interface::msg::ChassisServo>::SharedPtr
        controlSub;

    std::unique_ptr<CtrlAlg::PIDPos<double>> gyroPID;
    std::unique_ptr<CtrlAlg::PIDPos<double>> anglePID;
    std::unique_ptr<CtrlAlg::PIDPos<double>> speedPID;
    std::unique_ptr<CtrlAlg::PIDPos<double>> rotatePID;
    rclcpp::TimerBase::SharedPtr timer_balance;
    rclcpp::TimerBase::SharedPtr timer_angle;
    rclcpp::TimerBase::SharedPtr timer_speed;
    rclcpp::TimerBase::SharedPtr timer_publish_state;
    OnSetParametersCallbackHandle::SharedPtr param_callback_handle;

    std::map<std::string, rclcpp_lifecycle::LifecyclePublisher<
                              std_msgs::msg::Float64>::SharedPtr>
        publishers;

    struct State
    {
        enum class Target
        {
            Angle,
            Speed
        };
        double gyroy = 0;
        double gyroz = 0;
        double yaw = 0;
        double pitch = 0;
        double roll = 0;
        double speed = 0;

        double forwardTorque = 0;
        double rotateTorque = 0;
        double leftCurrent = 0;
        double rightCurrent = 0;

        double leftVec = 0;
        double rightVec = 0;

        double targetGyro = 0;
        double targetAngle = 0;
        double targetSpeed = 0;
        double targetRotate = 0;
        Target target = Target::Speed;
    };
    State state;
    double maxGyro;
    double maxAngle;
    double maxSpeed;
    double angle_offset;
    double gyro_angle_forward;

    std::unique_ptr<DataFilter> gyro_df;
    std::unique_ptr<DataFilter> angle_df;
    std::unique_ptr<DataFilter> speed_df;
    std::unique_ptr<DataFilter> gyroz_df;

  public:
    ControlNode() : rclcpp_lifecycle::LifecycleNode("control_node")
    {
        RCLCPP_INFO(get_logger(), "load parameters for PID");
        gyroPID =
            CreatPIDWithParam<CtrlAlg::PIDPos<double>>("gyro", 0.1, 0, 0, 1000);
        anglePID = CreatPIDWithParam<CtrlAlg::PIDPos<double>>("angle", 5.0, 0.1,
                                                              0, 1000);
        speedPID = CreatPIDWithParam<CtrlAlg::PIDPos<double>>("speed", 0.2, 1.0,
                                                              0, 1000);
        rotatePID = CreatPIDWithParam<CtrlAlg::PIDPos<double>>("rotate", 0.2,
                                                               1.0, 0, 1000);
        declare_parameter<float>("max_gyro", 30.0);
        maxGyro = get_parameter("max_gyro").as_double();
        declare_parameter<float>("max_angle", 10.0);
        maxAngle = get_parameter("max_angle").as_double();
        declare_parameter<float>("max_speed", 1.0);
        maxSpeed = get_parameter("max_speed").as_double();
        declare_parameter<float>("angle_offset", 0.0);
        angle_offset = get_parameter("angle_offset").as_double();
        declare_parameter<float>("gyro_angle_forward", 0.2);
        gyro_angle_forward = get_parameter("gyro_angle_forward").as_double();

        declare_parameter<float>("gyro_filter", 0.5);
        double gyro_filter = get_parameter("gyro_filter").as_double();
        declare_parameter<float>("angle_filter", 0.5);
        double angle_filter = get_parameter("angle_filter").as_double();
        declare_parameter<float>("speed_filter", 0.5);
        double speed_filter = get_parameter("speed_filter").as_double();
        declare_parameter<float>("gyroz_filter", 0.5);
        double gyroz_filter = get_parameter("gyroz_filter").as_double();
        gyro_df = std::make_unique<DataFilter>(gyro_filter, 20);
        angle_df = std::make_unique<DataFilter>(angle_filter, 20);
        speed_df = std::make_unique<DataFilter>(speed_filter, 20);
        gyroz_df = std::make_unique<DataFilter>(gyroz_filter, 5);

        RCLCPP_INFO(get_logger(), "PID parameters loaded");

        declare_parameter<double>("state_rate", -1.0);

        param_callback_handle = RegeistParamCallBack();
    }

  private:
    template <class T>
    std::unique_ptr<T> CreatPIDWithParam(std::string name, double p, double i,
                                         double d, double maxIOut)
    {
        std::vector<double> values = {p, i, d, maxIOut};
        rcl_interfaces::msg::ParameterDescriptor descriptor;
        descriptor.description = "PID parameter for " + name;
        declare_parameter<std::vector<double>>(name + "_pid", values,
                                               descriptor);
        values = get_parameter(name + "_pid").as_double_array();
        if (values.size() != 4)
        {
            RCLCPP_ERROR(get_logger(), "Invalid parameter size for %s",
                         name.c_str());
            return nullptr;
        }
        RCLCPP_INFO(get_logger(), "PID parameter for %s: %f, %f, %f",
                    name.c_str(), values[0], values[1], values[2]);
        auto pid = std::make_unique<CtrlAlg::PIDPos<double>>(
            values[0], values[1], values[2]);
        pid->SetMaxIntegralOutputPart(values[3]);
        return pid;
    }
    OnSetParametersCallbackHandle::SharedPtr RegeistParamCallBack()
    {
        return add_on_set_parameters_callback(
            [&](const std::vector<rclcpp::Parameter> &parameters)
            {
                rcl_interfaces::msg::SetParametersResult result;
                result.successful = true;
                for (const auto &param : parameters)
                {
                    if (param.get_name() == "gyro_pid")
                    {
                        auto values = param.as_double_array();
                        gyroPID->SetPID(values[0], values[1], values[2]);
                        RCLCPP_INFO(get_logger(), "gyro PID: %f, %f, %f",
                                    values[0], values[1], values[2]);
                    }
                    else if (param.get_name() == "angle_pid")
                    {
                        auto values = param.as_double_array();
                        anglePID->SetPID(values[0], values[1], values[2]);
                        RCLCPP_INFO(get_logger(), "angle PID: %f, %f, %f",
                                    values[0], values[1], values[2]);
                    }
                    else if (param.get_name() == "speed_pid")
                    {
                        auto values = param.as_double_array();
                        speedPID->SetPID(values[0], values[1], values[2]);
                        RCLCPP_INFO(get_logger(), "speed PID: %f, %f, %f",
                                    values[0], values[1], values[2]);
                    }
                    else if (param.get_name() == "rotate_pid")
                    {
                        auto values = param.as_double_array();
                        rotatePID->SetPID(values[0], values[1], values[2]);
                        RCLCPP_INFO(get_logger(), "rotate PID: %f, %f, %f",
                                    values[0], values[1], values[2]);
                    }
                    else if (param.get_name() == "state_rate")
                    {
                        result.successful = true;
                    }
                    else
                    {
                        result.successful = false;
                        result.reason =
                            "unknown parameter name: " + param.get_name();
                        break;
                    }
                }
                return result;
            });
    }

    using CallbackReturn = rclcpp_lifecycle::node_interfaces::
        LifecycleNodeInterface::CallbackReturn;
    CallbackReturn on_configure(const rclcpp_lifecycle::State &) override
    {
        RCLCPP_INFO(get_logger(), "Creat ROS2 related stuff");
        leftSpeedSub = create_subscription<WheelSpeed>(
            "left_speed", 1,
            [this](WheelSpeed::SharedPtr msg) { state.leftVec = msg->data; });
        rightSpeedSub = create_subscription<WheelSpeed>(
            "right_speed", 1,
            [this](WheelSpeed::SharedPtr msg) { state.rightVec = msg->data; });
        leftTorquePub = create_publisher<WheelTorque>("left_torque", 1);
        rightTorquePub = create_publisher<WheelTorque>("right_torque", 1);
        imuSub = create_subscription<sensor_msgs::msg::Imu>(
            "imu", 1,
            [this](sensor_msgs::msg::Imu::SharedPtr msg)
            {
                state.gyroy =
                    msg->angular_velocity.y * std::numbers::inv_pi * 180;
                state.gyroz =
                    msg->angular_velocity.z * std::numbers::inv_pi * 180;
                tf2::Quaternion q(msg->orientation.x, msg->orientation.y,
                                  msg->orientation.z, msg->orientation.w);
                tf2::Matrix3x3 m(q);
                m.getRPY(state.roll, state.pitch, state.yaw);
                state.roll *= std::numbers::inv_pi * 180;
                state.pitch *= std::numbers::inv_pi * 180;
                state.yaw *= std::numbers::inv_pi * 180;

                state.gyroy = gyro_df->filter(state.gyroy);
                state.gyroz = gyroz_df->filter(state.gyroz);
                state.pitch = angle_df->filter(state.pitch);
            });
        controlSub = create_subscription<general_interface::msg::ChassisServo>(
            "control_signal", 1,
            [this](general_interface::msg::ChassisServo::SharedPtr msg)
            {
                if (msg->angle_gain > 0.9)
                {
                    state.targetAngle =
                        std::clamp(msg->angle * std::numbers::inv_pi * 180,
                                   -maxAngle, maxAngle);
                    state.target = State::Target::Angle;
                }
                else if (msg->speed_gain > 0.9)
                {
                    state.targetSpeed =
                        std::clamp(msg->speed, -maxSpeed, maxSpeed);
                    state.target = State::Target::Speed;
                }
                state.targetRotate = std::clamp(
                    msg->gyroz * std::numbers::inv_pi * 180, -maxGyro, maxGyro);
            });
        RCLCPP_INFO(get_logger(), "Finish creating ROS2 related stuff");

        RCLCPP_INFO(get_logger(), "creat timer for control loop");
        timer_balance =
            create_wall_timer(std::chrono::milliseconds(1000 / gyro_rate),
                              std::bind(&ControlNode::OnGyroLoop, this));
        timer_angle =
            create_wall_timer(std::chrono::milliseconds(1000 / angle_rate),
                              std::bind(&ControlNode::OnAngleLoop, this));
        timer_speed =
            create_wall_timer(std::chrono::milliseconds(1000 / speed_rate),
                              std::bind(&ControlNode::OnSpeedLoop, this));
        timer_balance->cancel();
        timer_angle->cancel();
        timer_speed->cancel();
        double state_publish_rate = get_parameter("state_rate").as_double();
        if (state_publish_rate > 0)
        {
            timer_publish_state = create_wall_timer(
                std::chrono::milliseconds((int)(1000 / state_publish_rate)),
                std::bind(&ControlNode::OnPublishState, this));
            timer_publish_state->cancel();
        }
        RCLCPP_INFO(get_logger(), "timer for control loop created");
        RCLCPP_INFO(get_logger(), "tri-pid configure finished");
        return CallbackReturn::SUCCESS;
    }

    CallbackReturn on_activate(const rclcpp_lifecycle::State &) override
    {
        timer_balance->reset();
        timer_angle->reset();
        timer_speed->reset();
        if (timer_publish_state)
            timer_publish_state->reset();

        leftTorquePub->on_activate();
        rightTorquePub->on_activate();
        for (auto [name, pub] : publishers)
        {
            pub->on_activate();
        }
        RCLCPP_INFO(get_logger(), "tri-pid activate finished");
        return CallbackReturn::SUCCESS;
    }
    CallbackReturn on_deactivate(const rclcpp_lifecycle::State &) override
    {
        timer_balance->cancel();
        timer_angle->cancel();
        timer_speed->cancel();
        if (timer_publish_state)
            timer_publish_state->cancel();
        std_msgs::msg::Float64 msg;
        msg.data = 0;
        leftTorquePub->publish(msg);
        rightTorquePub->publish(msg);
        leftTorquePub->on_deactivate();
        rightTorquePub->on_deactivate();
        for (auto [name, pub] : publishers)
        {
            pub->on_deactivate();
        }
        RCLCPP_INFO(get_logger(), "tri-pid deactivate finished");
        return CallbackReturn::SUCCESS;
    }
    CallbackReturn on_cleanup(const rclcpp_lifecycle::State &) override
    {
        RCLCPP_INFO(get_logger(), "tri-pid cleanup finished");
        return CallbackReturn::SUCCESS;
    }

    void OnGyroLoop()
    {
        state.forwardTorque = -gyroPID->GetControl(
            state.targetGyro - state.gyroy, 1.0 / gyro_rate);
        state.forwardTorque +=
            gyro_angle_forward * (state.pitch - angle_offset);
        state.forwardTorque = ClampAbs(state.forwardTorque, 50);

        state.rotateTorque = -rotatePID->GetControl(
            state.targetRotate - state.gyroz, 1.0 / gyro_rate);
        state.rotateTorque = ClampAbs(state.rotateTorque, 50);

        state.leftCurrent =
            ClampAbs(state.forwardTorque + state.rotateTorque, 80);
        state.rightCurrent =
            ClampAbs(state.forwardTorque - state.rotateTorque, 80);

        WheelSpeed msg;
        msg.data = state.leftCurrent;
        leftTorquePub->publish(msg);
        msg.data = state.rightCurrent;
        rightTorquePub->publish(msg);
    }

    void OnAngleLoop()
    {
        state.targetGyro = anglePID->GetControl(
            state.targetAngle - angle_offset - state.pitch, 1.0 / angle_rate);
        state.targetGyro = ClampAbs(state.targetGyro, maxGyro);
        state.speed = (state.leftVec + state.rightVec) / 2;
        state.speed = speed_df->filter(state.speed);
    }

    void OnSpeedLoop()
    {
        if (state.target == State::Target::Speed)
        {
            state.targetAngle = speedPID->GetControl(
                state.targetSpeed - state.speed, 1.0 / speed_rate);
            state.targetAngle = ClampAbs(state.targetAngle, maxAngle);
        }
    }

    void OnPublishState()
    {
        Publish("torque", state.forwardTorque);
        Publish("rotate_torque", state.rotateTorque);
        Publish("target_angle", state.targetAngle);
        Publish("angle", state.pitch);
        Publish("target_gyro", state.targetGyro);
        Publish("gyro", state.gyroy);
        Publish("target_gyroz", state.targetRotate);
        Publish("gyroz", state.gyroz);
        Publish("target_speed", state.targetSpeed);
        Publish("speed", state.speed);
    }

    void Publish(std::string topic, double value)
    {
        if (publishers.find(topic) == publishers.end())
            publishers[topic] =
                create_publisher<std_msgs::msg::Float64>(topic, 1);
        auto msg = std_msgs::msg::Float64();
        msg.data = value;
        publishers[topic]->publish(msg);
    }

    double ClampAbs(double value, double max)
    {
        if (value > max)
        {
            return max;
        }
        else if (value < -max)
        {
            return -max;
        }
        return value;
    }

    bool InRange(double value, double min, double max)
    {
        return value >= min && value <= max;
    }
};

int main(int argc, char **argv)
{
    rclcpp::init(argc, argv);
    auto node = std::make_shared<ControlNode>();
    node->trigger_transition(
        lifecycle_msgs::msg::Transition::TRANSITION_CONFIGURE);
    rclcpp::spin(node->get_node_base_interface());
    rclcpp::shutdown();
    return 0;
}